/**
 * @file main.cpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-10-26
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */
#include <lyn_api.h>
#include <fstream>
#include <iostream>
#include <nlohmann/json.hpp>
#include <vector>
#include "argsParser.h"
#include "imageDecoder.h"
#include "imageEncoder.h"
#include "ipeYolov5Param.h"
#include "opencvWindow.hpp"
#include "osd_plugin.h"
#include "util.hpp"
#include "yolox_post_process.h"

using namespace std;

std::string GetBasePath(const std::string &filePath) {
  size_t lastSeparatorPos = filePath.find_last_of("/\\");
  if (lastSeparatorPos != std::string::npos) {
    return filePath.substr(0, lastSeparatorPos);
  }
  return "";
}

int main(int argc, char *argv[]) {
  // 1. 命令行解析
  if (ArgsParse::yolov5ArgsParser(argc, argv) == false) {
    return -1;
  }

  std::string strModelPath = ArgsParse::argModelPath;
  std::string strJpegPath = ArgsParse::argInputFilePath;

  // 2. 创建 context 与 stream
  lynContext_t ctx = nullptr;
  CHECK_ERR(lynCreateContext(&ctx, ArgsParse::argChipsVec[0]));
  CHECK_ERR(lynRegisterErrorHandler(StreamErrorHandler, nullptr));
  lynStream_t stream = nullptr;
  CHECK_ERR(lynCreateStream(&stream));

  std::string strJpegOutPath = strJpegPath;
  strJpegOutPath.insert(strJpegPath.rfind('.'), "_result");

  std::string label_file = "/../yolox/label/label.v5x7.json";
  std::string label_path = GetBasePath(ArgsParse::argPostPluginPath) + label_file;

  ifstream fJson(label_path);
  stringstream ss;
  ss << fJson.rdbuf();
  nlohmann::json js = nlohmann::json::parse(ss.str());

  int labelCount = js.size();
  LabelList labelList;

  if (labelCount == 0) {
    cout << "error : label file is not invalid." << endl;
    return 0;
  }

  labelList.labelNum = labelCount;
  labelList.labels = new LabelInfo[labelCount];
  memset(labelList.labels, 0, sizeof(LabelInfo) * labelCount);
  string labelName;
  for (int i = 0; i < labelCount; ++i) {
    labelName = js[i]["desc"].get<string>();
    strncpy(labelList.labels[i].name, labelName.c_str(),
            sizeof(labelList.labels[i].name) - 1);
  }

  // 3. 初始化解码图片类
  ImageDecoder imageDec;
  imageDec.Init(strJpegPath);
  lynImageInfo_t imgInfo;
  imageDec.GetImageDecInfo(&imgInfo);

  // 4. 同步解码图片到 Device 侧
  lynCodecBuf_t decodeImg;
  decodeImg.size = imgInfo.output.predictBufSize;
  CHECK_ERR(lynMalloc((void **)&decodeImg.data, decodeImg.size));
  imageDec.DecodeImageToDevice(stream, decodeImg.data);

  // 5. 加载模型
  ModelInfo yolov5Info;
  yolov5Info.LoadModelByPath(strModelPath.c_str());
  lynPixelFormat_t pixelFormat = LYN_PIX_FMT_NV12;

  // 6. 加载 Plugin, 并设置后处理参数
  lynPlugin_t postplugin;
  CHECK_ERR(lynPluginRegister(&postplugin, ArgsParse::argPostPluginPath.c_str()));
  lynPlugin_t osdplugin;
  CHECK_ERR(lynPluginRegister(&osdplugin, ArgsParse::argOsdPluginPath.c_str()));

  YoloxPostProcessInfo_t post_info;
  memset(&post_info, 0, sizeof(YoloxPostProcessInfo_t));
  post_info.is_pad_resize = 1;
  post_info.score_threshold = 0.25;
  post_info.nms_threshold = 0.45;
  post_info.nms_top_k = 500;
  post_info.anchorSize = 3;
  post_info.width = yolov5Info.width;
  post_info.height = yolov5Info.height;
  post_info.ori_width = imgInfo.output.width;
  post_info.ori_height = imgInfo.output.height;
  post_info.class_num = labelCount;
  LabelInfo *devLabels = nullptr;
  CHECK_ERR(
      lynMalloc((void **)&devLabels, sizeof(LabelInfo) * labelList.labelNum));
  CHECK_ERR(lynMemcpy(devLabels, labelList.labels,
                      sizeof(LabelInfo) * labelList.labelNum, ClientToServer));
  post_info.labelList.labelNum = labelList.labelNum;
  post_info.labelList.labels = devLabels;
  lynBoxesInfo *pDevBoxesInfo;
  CHECK_ERR(lynMalloc((void **)&pDevBoxesInfo, sizeof(lynBoxesInfo)));
  post_info.boxesInfo = (lynBoxesInfo *)pDevBoxesInfo;

  // 7. IPE 处理
  uint8_t *pIpeBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pIpeBufOut, yolov5Info.inputSize));
  IpeParamYolov5 ipeParam(yolov5Info.width, yolov5Info.height);
  ipeParam.SetImgInfo(imgInfo.output.width, imgInfo.output.height, pixelFormat);
  ipeParam.CalcParam(stream, decodeImg.data, pIpeBufOut);
  // 8. 调用 APU 推理接口，之后调用 Plugin 进行后处理
  uint8_t *pDevApuOutData = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDevApuOutData, yolov5Info.outputSize));
  CHECK_ERR(lynExecuteModelAsync(stream, yolov5Info.model, pIpeBufOut,
                                 pDevApuOutData, yolov5Info.batchSize));
  post_info.output_tensor = pDevApuOutData;
  CHECK_ERR(lynPluginRunAsync(stream, postplugin, "lynYoloxPostProcess",
                              &post_info, sizeof(post_info)));
  lynSynchronizeStream(stream);
  lynDrawBoxAndTextPara para;
  para.imgData = decodeImg.data;
  para.imgFmt = LYN_PIX_FMT_NV12;
  para.imgW = imgInfo.output.width;
  para.imgH = imgInfo.output.height;
  para.boxesInfo = pDevBoxesInfo;
  para.boxColor = DRAW_COLOR_BLUE;
  para.boxThick = DRAW_THICK_2;
  para.fontSize = FONT_SIZE_24;
  para.fontColor = DRAW_COLOR_BLUE;
  CHECK_ERR(lynPluginRunAsync(stream, osdplugin, "lynDrawBoxAndText", &para,
                              sizeof(para)));
  lynSynchronizeStream(stream);
  // 等待 plugin 处理完，释放内存
  CHECK_ERR(lynFree(pDevApuOutData));
  CHECK_ERR(lynFree(pIpeBufOut));

  // 9. 根据传入的 ShowType 决定是进行图片编码还是直接显示
  if (ArgsParse::argShowType == ShowType::SaveFile) {
    ImageEncoder encoder;
    encoder.EncodeImageAndSave(stream, strJpegOutPath, decodeImg, imgInfo);
  } else if (ArgsParse::argShowType == ShowType::DirectShow) {
    std::string windowName = "inferImageYolov5";
    int showTime = 5000;
    OpencvWindowManager::GetInstance()->ShowImage(
        windowName, &decodeImg, imgInfo.output.width, imgInfo.output.height,
        showTime);
  }
  // 10. 销毁资源
  CHECK_ERR(lynDestroyStream(stream));
  CHECK_ERR(lynFree(decodeImg.data));
  yolov5Info.UnLoadModel();
  if (ctx) {
    CHECK_ERR(lynDestroyContext(ctx));
  }

  return 0;
}
